import torch
from pathlib import Path
import h5py
import numpy as np

from typing import Union, Tuple
from torch import Tensor


class HDF5Dataset(torch.utils.data.Dataset):

    def __init__(self, file_path: Union[str, Path]):
        super(HDF5Dataset, self).__init__()
        self.file = h5py.File(file_path, 'r')
        self.X_ecg = self.file['X']
        self.X_cwt = self.file['X_cwt']
        self.rr = self.file['rr']
        self.y = self.file['y']
    

    def __del__(self):
        self.close()


    def close(self):
        self.file.close()


    def __getitem__(self, index: Union[int, slice, np.ndarray, list]) -> Tuple[Tensor, Tensor, Tensor]:
        if isinstance(index, int) or isinstance(index, slice):
            X_ecg_data = torch.tensor(self.X_ecg[index], dtype=torch.float32)
            X_cwt_data = torch.tensor(self.X_cwt[index], dtype=torch.float32)
            rr_data = torch.tensor(self.rr[index], dtype=torch.float32)
            y_data = torch.tensor(self.y[index], dtype=torch.long)
        else:
            # Fetch one item at a time
            X_ecg_data = torch.stack([torch.tensor(self.X_ecg[i], dtype=torch.float32) for i in index])
            X_cwt_data = torch.stack([torch.tensor(self.X_cwt[i], dtype=torch.float32) for i in index])
            rr_data = torch.stack([torch.tensor(self.rr[i], dtype=torch.float32) for i in index])
            y_data = torch.stack([torch.tensor(self.y[i], dtype=torch.long) for i in index])
            
        return X_ecg_data, X_cwt_data, rr_data, y_data


    def __len__(self):
        return len(self.y)
